package com.kodin.polyfit;

import com.blankj.utilcode.util.LogUtils;

import org.apache.commons.math3.stat.regression.RegressionResults;
import org.apache.commons.math3.stat.regression.SimpleRegression;

import java.util.ArrayList;
import java.util.List;

public class LinearFit {

    public static Result linearFit(double[][] data) {
        List<double[]> fitData = new ArrayList<>();
        SimpleRegression regression = new SimpleRegression();
        regression.addData(data); // 数据集
        /*
         * RegressionResults 中是拟合的结果
         * 其中重要的几个参数如下：
         *   parameters:
         *      0: b
         *      1: k
         *   globalFitInfo
         *      0: 平方误差之和, SSE
         *      1: 平方和, SST
         *      2: R 平方, RSQ
         *      3: 均方误差, MSE
         *      4: 调整后的 R 平方, adjRSQ
         *
         * */
        RegressionResults results = regression.regress();
        double b = results.getParameterEstimate(0);
        double k = results.getParameterEstimate(1);
        double[] best = {b, k};
        double r2 = results.getRSquared();
        String re = String.format("b:%f,k:%f,r2:%f\n", b, k, r2);
        LogUtils.e(re);

        // 重新计算生成拟合曲线
        for (double[] datum : data) {
            double[] xy = {datum[0], k * datum[0] + b};
            fitData.add(xy);
        }

        StringBuilder func = new StringBuilder();
        func.append(re);
        func.append("f(x) =");
        func.append(b >= 0 ? " " : " - ");
        func.append(Math.abs(b));
        func.append(k > 0 ? " + " : " - ");
        func.append(Math.abs(k));
        func.append("x");
        return new Result(data, fitData.stream().toArray(double[][]::new), func.toString(), best);
    }

    /**
     * y = kx + b
     * f(x) = 1.5x + 0.5
     *
     * @return
     */
    public static double[][] linearScatters() {
        List<double[]> data = new ArrayList<>();
        for (double x = 0; x <= 10; x += 1) {
            double y = 1.5 * x + 0.5;
//            y += Math.random() * 2 - 1; // 随机数
            double[] xy = {x, y};
            data.add(xy);
        }
        return data.stream().toArray(double[][]::new);

    }

}
